import os
from http import HTTPStatus
from pathlib import PurePosixPath
from urllib.parse import urlparse, unquote

import dashscope
import requests
from dashscope import ImageSynthesis

default_max_retry_count = 3

dashscope.api_key = 'sk-9a53285a36204cf589bc4a443a78f534'
model = "flux-schnell"

def sample_block_call(input_prompt, file_path):
    rsp = ImageSynthesis.call(model=model,
                              prompt=input_prompt,
                              size='1024*1024')
    if rsp.status_code == HTTPStatus.OK:
        print(rsp.output)
        print(rsp.usage)

        if rsp.output.task_status == "FAILED":
            return {"success": False, "value": None, "msg": "生成图片失败"}

        if not os.path.exists(file_path):
            os.makedirs(file_path)

        for result in rsp.output.results:
            file_name = PurePosixPath(unquote(urlparse(result.url).path)).parts[-1]
            dest_file_path = ""
            if file_path is not None:
                file_name = file_path + os.sep + file_name

            with open('./%s' % file_name, 'wb+') as f:
                f.write(requests.get(result.url).content)
        return {"success": True, "value": file_name, "msg": "生成图片成功"}
    else:
        print('Failed, status_code: %s, code: %s, message: %s' %
              (rsp.status_code, rsp.code, rsp.message))
        return {"success": False, "value": None, "msg": "生成图片失败"}

def sample_block_call_with_retry(input_prompt, file_path, max_retry_count):
    try:
        file_name = ""
        if max_retry_count is None or max_retry_count <= 0:
            max_retry_count = default_max_retry_count

        rsp = ImageSynthesis.call(model=model,
                                  prompt=input_prompt,
                                  size='1024*1024')
        if rsp.status_code != HTTPStatus.OK:
            retry_count = 1
            while(retry_count <= max_retry_count):
                print("有失败图片，重试第" + str(retry_count) + "次！")
                rsp = ImageSynthesis.call(model=model,
                                          prompt=input_prompt,
                                          size='1024*1024')
                if rsp.status_code == HTTPStatus.OK :
                    break
                retry_count = retry_count + 1

            if retry_count >= max_retry_count and rsp.status_code != HTTPStatus.OK:
                return {"success": False, "value": None, "msg": "生成图片失败"}

        if rsp.output.task_status == "FAILED":
            retry_count = 1
            while (retry_count <= max_retry_count):
                print("有失败图片，重试第" + str(retry_count) + "次！")
                rsp = ImageSynthesis.call(model=model,
                                          prompt=input_prompt,
                                          size='1024*1024')
                if rsp.output.task_status == "SUCCESS":
                    break
                retry_count = retry_count + 1

            if retry_count >= max_retry_count and rsp.output.task_status == "FAILED":
                return {"success": False, "value": None, "msg": "生成图片失败"}

        if not os.path.exists(file_path):
            os.makedirs(file_path)

        for result in rsp.output.results:
            file_name = PurePosixPath(unquote(urlparse(result.url).path)).parts[-1]
            if file_name is None or file_name == "":
                continue

            if file_path is not None:
                file_name = file_path + os.sep + file_name

            with open('./%s' % file_name, 'wb+') as f:
                f.write(requests.get(result.url).content)

        return {"success": True, "value": file_name, "msg": "生成图片成功"}
    except Exception as e:
        print(e)
        return {"success": False, "value": None, "msg": "生成图片异常,信息为:" + str(e)}

def sample_async_call(input_prompt):
    rsp = ImageSynthesis.async_call(model=model,
                                    prompt=input_prompt,
                                    size='1024*1024')
    if rsp.status_code == HTTPStatus.OK:
        print(rsp.output)
        print(rsp.usage)
    else:
        print('Failed, status_code: %s, code: %s, message: %s' %
              (rsp.status_code, rsp.code, rsp.message))
    status = ImageSynthesis.fetch(rsp)
    if status.status_code == HTTPStatus.OK:
        print(status.output.task_status)
    else:
        print('Failed, status_code: %s, code: %s, message: %s' %
              (status.status_code, status.code, status.message))

    rsp = ImageSynthesis.wait(rsp)
    if rsp.status_code == HTTPStatus.OK:
        print(rsp.output)
    else:
        print('Failed, status_code: %s, code: %s, message: %s' %
              (rsp.status_code, rsp.code, rsp.message))
